from gym_pybullet_drones.envs.HoverAviary import HoverAviary
from gym_pybullet_drones.utils.enums import ObservationType, ActionType
'''
Total reward:  33542.44341415363
mean reward:  335.4244341415363
total step:  21817
mean step:  218.17
'''

import gymnasium as gym

DEFAULT_OBS = ObservationType('kin') # 'kin' or 'rgb'
DEFAULT_ACT = ActionType('one_d_rpm') # 'rpm' or 'pid' or 'vel' or 'one_d_rpm' or 'one_d_pid'
# env = HoverAviary(gui=True)
env = HoverAviary(obs=DEFAULT_OBS, act=DEFAULT_ACT)
# env = HoverAviary()

# Reset the environment to get the initial state
total_reward = 0
total_step = 0
# Run a loop to play the game
episoid = 100
for _ in range(episoid):
    state = env.reset()
    while True:
        # Take a random action
        # env.render()
        action = env.action_space.sample()

        # Get the next state, reward, done flag, and info from the environment
        state, reward, done, trunc, info = env.step(action)
        total_step += 1
        if reward != 0:
            total_reward += reward
            print("action: ", action)   
            print("reward: ", reward)
            print("info: ", info)

        # If done, reset the environment
        if done or trunc:
            print(f"done: {done} trunc: {trunc}")
            print(f"action: ", action)
            print("reward: ", reward)
            print("info: ", info)
        #     print("count_frame: ", count_frame)``
            break

print("Total reward: ", total_reward)
print("mean reward: ", total_reward / episoid)
print("total step: ", total_step)
print("mean step: ", total_step / episoid)

# Close the environment
env.close()
